Basic CUDA support #68
Conversation
|
|
||
| # CPU idiosyncracy that needs to be done differently | ||
| coord_ramps = torch.from_numpy(np.mgrid[slices]) | ||
| coord_ramps = torch.from_numpy(np.mgrid[slices]).to(points.device) |
There was a problem hiding this comment.
this will work, but allocates an array on cpu, then sends it to gpu. we may want to borrow from the prior code that uses torch.meshgrid(x_vals, y_vals, z_vals) after allocating x_vals=torch.arange(start, end, device=device) etc. that way it gets created on gpu.
This is second-order optimization probably, since there will likely be other bottlenecks to fix beforehand, so keep as is for now
|
|
||
| grads_points = None | ||
| grad_values = None | ||
|
|
||
| ndim = points.shape[0] | ||
|
|
||
| nufft_func = get_nufft_func(ndim, 2) | ||
| nufft_func = get_nufft_func(ndim, 2, points.device.type) |
There was a problem hiding this comment.
Sending points.device object doesn't work?
BTW do we know anything about how well cufinufft interacts with multple devices?
There was a problem hiding this comment.
I'm guessing cufinufft does not like multiple devices, but I haven't tried.
We definitely need more checks that the arrays are both on the same device (at least cpu/cuda, if not also checking they're on the same index of cuda)
There was a problem hiding this comment.
Oh and we could use points.device, but the only thing we care about for now is if it is cuda/cpu, so sending the type seemed simplest
Based off of #64.
Locally, the added type1 cuda tests are passing.
CI will probably take a couple iterations before it is all working.
There's not really much to this: I extended the
get_nufft_funchelper and switched to manually handling the necessary fftshifts, since the cuda functions don't (currently?) acceptmodeord